%%%-------------------------------------------------------------------
%%%    BASIC INFORMATION
%%%-------------------------------------------------------------------
%%% @copyright 2006 Erlang Training & Consulting Ltd
%%% @author  Martin Carlson <martin@erlang-consulting.com>
%%% @version 0.0.1
%%% @doc
%%% @end
%%%-------------------------------------------------------------------
-module(psql_logic).

-behaviour(gen_fsm).

%% API
-export([start_link/0,
	 start_link/1,
	 connect/6,
	 disconnect/1,
	 q/2]).

%% Internal API
-export([connection_event/2,
	 connection_closed/1]).

%% States
-export([disconnected/3,
	 disconnected/2,
	 authentication/2,
	 authentication/3,
	 setup/2,
	 setup/3,
	 connected/3,
	 connected/2,
	 command_response/2,
	 command_response/3,
	 row_data/2,
	 row_data/3,
	 copy/2,
	 copy/3]).

%% gen_fsm callbacks
-export([init/1, 
	 handle_event/3,
	 handle_sync_event/4, 
	 handle_info/3, 
	 terminate/3, 
	 code_change/4]).

-record(state, {pool_name = default,
		connection,
		user,
		password,
		db,
		digest,
		from,
		row_description,
		row_acc = []}).

-include("psql.hrl").
-define(PROTOCOL, 196608).

%%====================================================================
%% API
%%====================================================================
start_link() ->
    {ok, Pid} = gen_fsm:start_link(?MODULE, [default], []),
    case application:get_env(psql, hostspec) of
	{ok, {Host, Port, Usr, Pwd, Db}} ->
	    ok = connect(Pid, Host, Port, Usr, Pwd, Db),
	    {ok, Pid};
	_Reason ->
	    {ok, Pid}
    end.

start_link(PoolName) ->
    {ok, Pid} = gen_fsm:start_link(?MODULE, [PoolName], []),
    case application:get_env(psql, PoolName) of
	{ok, {Host, Port, Usr, Pwd, Db}} ->
	    ok = connect(Pid, Host, Port, Usr, Pwd, Db),
	    {ok, Pid};
	_Reason ->
	    {ok, Pid}
    end.

connect(Pid, Host, Port, Usr, Pwd, Db) ->
    gen_fsm:sync_send_event(Pid, {connect, Host, Port, Usr, Pwd, Db}).

disconnect(Pid) ->
    gen_fsm:sync_send_event(Pid, disconnect).

q(Pid, Query) ->
    gen_fsm:sync_send_event(Pid, {q, Query}, infinity).

%%====================================================================
%% API
%%===================================================================
connection_event(Pid, Event) ->
    gen_fsm:send_event(Pid, Event).

connection_closed(Pid) ->
    gen_fsm:send_all_state_event(Pid, connection_closed).

%%====================================================================
%% gen_fsm callbacks
%%====================================================================
init([PoolName]) ->
    process_flag(trap_exit, true),
    {ok, Connection} = psql_connection:start_link(self()),
    {ok, disconnected, #state{pool_name = PoolName,
			      connection = Connection}}.

%%====================================================================
%% Disconnected States
%%====================================================================
disconnected(connection_established, State) ->
    {AuthMsg, AuthDigest} = psql_protocol:authenticate(?PROTOCOL,
						       State#state.user,
						       State#state.password,
						       State#state.db),
    psql_connection:command(State#state.connection, {send, AuthMsg}),
    {next_state, authentication, State#state{digest = AuthDigest}};
disconnected(_Event, State) ->
    {next_state, disconnected, State}.

disconnected({connect, Host, Port, Usr, Pwd, Db}, _From, State) ->
    psql_connection:command(State#state.connection, {connect, Host, Port}),
    {reply, ok, disconnected, State#state{user = Usr, password = Pwd, db = Db}};
disconnected({q, _}, _From, State) ->
    {reply, {error, disconnected}, disconnected, State};
disconnected(_Event, _From, State) ->
    {reply, ok, disconnected, State}.

%%====================================================================
%% Authentication States
%%====================================================================
authentication(connection_closed, State) ->
    {next_state, disconnected, State};
authentication({psql, authentication, <<0,0,0,5, Salt/binary>>}, State) ->
    AuthDigest = psql_protocol:md5digest(State#state.digest, Salt),
    psql_connection:command(State#state.connection, {send, AuthDigest}),
    {next_state, authentication, State};
authentication({psql, authentication, <<0,0,0,0>>}, State) ->
    {next_state, setup, State};
authentication({psql, error, Error}, State) ->
    case psql_protocol:error(Error) of
	{authentication, _} ->	    
	    psql_connection:command(State#state.connection, disconnect),
	    error_logger:error_report({psql, authentication_failed}),
	    {next_state, disconnected, State};
	{_, _} ->
	    {next_state, disconnected, State}
    end;
authentication(Event, State) ->
    io:format("Event: ~p~n", [Event]),
    {next_state, authentication, State}.

authentication({q, _}, _From, State) ->
    {reply, {error, disconnected}, authentication, State};
authentication(Event, _From, State) ->
    io:format("Event: ~p~n", [Event]),
    {reply, ok, authentication, State}.

%%====================================================================
%% Setup States
%%====================================================================
setup({psql,ready_for_query, _}, State) ->
    sql:register_resource(self(), psql, State#state.pool_name),
    {next_state, connected, State};
setup({error, Error}, State) ->
    case psql_protocol:error(Error) of
	{shutdown, _} ->
	    error_logger:error_report({psql, shutdown}),
	    {next_state, disconnected, State};
	{_, _} ->
	    {next_state, disconnected, State}
    end;
setup(Event, State) ->
    io:format("Event: ~p~n", [Event]),
    {next_state, setup, State}.

setup({q, _}, _From, State) ->
    {reply, {error, disconnected}, setup, State};
setup(Event, _From, State) ->
    io:format("Event: ~p~n", [Event]),
    {reply, ok, setup, State}.

%%====================================================================
%% Connected States
%%====================================================================
connected({error, Error}, State) ->
    case psql_protocol:error(Error) of
	{shutdown, _} ->
	    error_logger:error_report({psql, shutdown}),
	    {next_state, disconnected, State};
	{_, _} ->
	    {next_state, disconnected, State}
    end;
connected(_Event, State) ->
    {next_state, connected, State}.

connected(disconnect, _From, State) ->
    psql_connection:command(State#state.connection, disconnect),
    {reply, ok, disconnected, State};
connected({q, Query}, From, State) ->
    QueryMsg = psql_protocol:q(Query),
    psql_connection:command(State#state.connection, {send, QueryMsg}),
    {next_state, command_response, State#state{from = From}};
connected(_Event, _From, State) ->
    {reply, ok, connected, State}.

%%====================================================================
%% Row Description States
%%====================================================================
command_response({psql, row_description, Data}, State) ->
    Desc = row_description(Data),
    {next_state, row_data, State#state{row_description = Desc}};
command_response({psql, copy_in_response, _Data}, State) ->
    Cmd = psql_protocol:copy_done(),    
    psql_connection:command(State#state.connection, {send, Cmd}),    
    {next_state, copy, State};
command_response({psql, copy_out_response, _Data}, State) ->
    Cmd = psql_protocol:copy_done(),    
    psql_connection:command(State#state.connection, {send, Cmd}),
    {next_state, copy, State};
command_response({psql, ready_for_query, _}, State) ->
    gen_fsm:reply(State#state.from, {ok, {no_result, []}}),
    {next_state, connected, State#state{row_acc = []}};
command_response({psql, error, Error}, State) ->
    case psql_protocol:error(Error) of
	{sql_error, Code, _Desc} ->
	    gen_fsm:reply(State#state.from, {sql_error, Code}),
	    {next_state, connected, State#state{row_acc = []}};
	{_, _} ->
	    {next_state, disconnected, State#state{row_acc = []}}
    end;
command_response({error, Error}, State) ->
    case psql_protocol:error(Error) of
	{shutdown, _} ->
	    error_logger:error_report({psql, shutdown}),
	    {next_state, disconnected, State#state{row_acc = []}};
	{sql_error, Code, _Desc} ->
	    gen_fsm:reply(State#state.from, {sql_error, Code}),
	    {next_state, connected, State#state{row_acc = []}};
	{_, _} ->
	    {next_state, disconnected, State#state{row_acc = []}}
    end;
command_response(Event, State) ->
    io:format("Event (command_response): ~p~n", [Event]),
    {next_state, command_response, State}.

command_response({q, _}, _From, State) ->
    {reply, {error, busy}, command_response, State};
command_response(Event, _From, State) ->
    io:format("Event: ~p~n", [Event]),
    {reply, ok, command_response, State}.

%%====================================================================
%% Row Data States
%%====================================================================
row_data({psql, data_row, Data}, State) ->
    NewAcc = [row(Data, State#state.row_description)|State#state.row_acc],
    {next_state, row_data, State#state{row_acc = NewAcc}};
row_data({psql, command_complete, _}, State) ->
    Reply = {State#state.row_description, State#state.row_acc},
    gen_fsm:reply(State#state.from, {ok, Reply}),
    {next_state, row_data, State#state{row_acc = []}};
row_data({psql,ready_for_query, _}, State) ->
    {next_state, connected, State};
row_data({error, Error}, State) ->
    case psql_protocol:error(Error) of
	{shutdown, _} ->
	    error_logger:error_report({psql, shutdown}),
	    {next_state, disconnected, State};
	{_, _} ->
	    {next_state, disconnected, State}
    end;
row_data(Event, State) ->
    io:format("Event: ~p~n", [Event]),
    {next_state, row_data, State}.

row_data({q, _}, _From, State) ->
    {reply, {error, busy}, row_data, State};
row_data(Event, _From, State) ->
    io:format("Event (row_data): ~p~n", [Event]),
    {reply, ok, row_data, State}.

%%====================================================================
%% Copy States
%%====================================================================
copy(Event, State) ->
    io:format("Event (copy): ~p~n", [Event]),
    {next_state, copy, State}.

copy({q, _}, _From, State) ->
    {reply, {error, busy}, copy, State};
copy(Event, _From, State) ->
    io:format("Event (copy): ~p~n", [Event]),
    {reply, ok, copy, State}.


%%====================================================================
%% gen_fsm callbacks
%%====================================================================
handle_event(connection_closed, _StateName, State) ->
    sql:unregister_resource(self(), State#state.pool_name),
    {next_state, disconnected, State#state{from = undefined,
					   row_description = [],
					   row_acc = []}};
handle_event(_Event, StateName, State) ->
    {next_state, StateName, State}.

handle_sync_event(_Event,  _From, StateName, StateData) ->
    {next_state,  StateName,  StateData}.

handle_info({'EXIT', Pid, Reason}, StateName, State) ->
    sql:unregister_resource(self(), State#state.pool_name),
    error_logger:error_report({exitsignal, Pid, Reason, StateName, State}),
    {stop, Reason, State};
handle_info(_Info, StateName, State) ->
    {next_state, StateName, State}.

terminate(_Reason, _StateName, State) ->
    sql:unregister_resource(self(), State#state.pool_name),
    psql_connection:command(State#state.connection, shutdown),
    ok.

code_change(_OldVsn, StateName, State, _Extra) ->
    {ok, StateName, State}.

%%--------------------------------------------------------------------
%%% Internal functions
%%--------------------------------------------------------------------
row_description(<<_:1/big-unit:16, Fields/binary>>) ->
    list_to_tuple(field_desc_parser(Fields, [])).

field_desc_parser(<<>>, Acc) ->
	lists:reverse(Acc);
field_desc_parser(Fields, Acc) ->
    {Name, <<Table:1/big-signed-unit:32,
	    Col:1/big-signed-unit:16,
	    Type:1/big-signed-unit:32,
	    Size:1/big-signed-unit:16,
	    _Mod:1/big-signed-unit:32,
	    Format:1/big-signed-unit:16,
	    Tail/binary>>} = field_name_parser(Fields, []),
    Field = #field{name = Name,
		   table_code = Table,
		   field_code = Col,
		   type = psql_protocol:data_type(Type),
		   max_length = Size,
		   format = Format},
    field_desc_parser(Tail, [Field|Acc]).

field_name_parser(<<0, Tail/binary>>, Acc) ->
    {lists:reverse(Acc), Tail};
field_name_parser(<<C:8, Tail/binary>>, Acc) ->
    field_name_parser(Tail, [C|Acc]).

row(<<_:1/big-signed-unit:16, Cols/binary>>, Desc) ->
    col_parser(Cols, Desc, 1, []).

col_parser(<<>>, _, _, Acc) ->
    list_to_tuple(lists:reverse(Acc));
col_parser(<<Length:1/big-signed-unit:32, Rest/binary>>, Desc, N, Acc) ->
    if
	Length == -1 ->
	    Value = [],
	    Tail = Rest;
	Length > size(Rest) ->
	    Value = [],
	    Tail = Rest,
	    erlang:error({badarg, Length, N, Rest});
	true ->
	    <<SQLValue:Length/binary, Tail/binary>> = Rest,
	    Value = psql_lib:type_cast(SQLValue, element(N, Desc))
    end,
    col_parser(Tail, Desc, N + 1, [Value|Acc]).

